# compare_results.py
import cv2
import numpy as np
import matplotlib.pyplot as plt
import os
from glob import glob


def compare_rural_urban_performance():
    """对比模型在农村和城市数据上的表现"""
    rural_pred_dir = 'results/test_predictions'
    urban_pred_dir = 'results/urban_predictions'

    # 分析预测结果差异
    rural_stats = analyze_predictions(rural_pred_dir)
    urban_stats = analyze_predictions(urban_pred_dir)

    print("=== 跨域性能对比 ===")
    print("农村测试集预测结果:")
    for class_name, percentage in rural_stats.items():
        print(f"  {class_name}: {percentage:.2f}%")

    print("城市测试集预测结果:")
    for class_name, percentage in urban_stats.items():
        print(f"  {class_name}: {percentage:.2f}%")

    # 计算性能变化
    print("性能变化 (%):")
    for class_name in rural_stats:
        rural_val = rural_stats[class_name]
        urban_val = urban_stats[class_name]
        change = ((urban_val - rural_val) / rural_val) * 100
        print(f"  {class_name}: {change:+.2f}%")


def analyze_predictions(pred_dir):
    """分析预测结果的类别分布"""
    color_map = {
        tuple([255, 0, 0]): '建筑',
        tuple([0, 255, 0]): '道路',
        tuple([0, 0, 255]): '水域',
        tuple([255, 255, 0]): '荒地',
        tuple([0, 255, 255]): '森林',
        tuple([255, 0, 255]): '农田',
        tuple([0, 0, 0]): '背景'
    }

    pred_files = glob(os.path.join(pred_dir, '*_pred.png'))
    class_pixels = {cls: 0 for cls in color_map.values()}
    total_pixels = 0

    for pred_file in pred_files[:10]:  # 抽样分析
        pred = cv2.imread(pred_file)
        pred = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)

        h, w, _ = pred.shape
        total_pixels += h * w

        # 统计每个类别的像素数
        for color, class_name in color_map.items():
            mask = np.all(pred == color, axis=-1)
            class_pixels[class_name] += np.sum(mask)

    # 计算百分比
    class_percentages = {}
    for class_name, count in class_pixels.items():
        class_percentages[class_name] = (count / total_pixels) * 100

    return class_percentages


def visualize_domain_shift():
    """可视化域偏移影响"""
    # 这里可以添加具体的可视化代码
    pass


if __name__ == '__main__':
    compare_rural_urban_performance()